import numpy as np
import matplotlib.pyplot as mp
import part1 as pa

#mp.clf()

#Parametres

G = 9.81
l = 1    #m
m = 1    #kg
step = 0.1 #s
#Modelisation du probleme du pendule simple

def func():

    
    
    def f(theta,t):

        return theta[1]

    def g(theta,t):

        return (-G/l)*np.sin(theta[0])
    
    return [f,g]


theta0 = func()


#Frequence
#Cette fonction calcule la frequence d'oscillation en fonction de l'angle initial du
#pendule. Sur le graphe, en rouge, la constante sqrt(G/l)/2*pi.


def freq(pas):

    i=0
    frequence = np.zeros ([3,pas/5])
    

    for j in np.arange(2,pas,5):

        X = pa.meth_step_n(np.array([(np.pi)/j,0.]),0.,3,0.05,theta0,pa.runge_kutta_step)[0]

        frequence[0][i] = np.abs(X[1][1]/X[1][0])
        frequence[1][i] = (np.pi)/j
        frequence[2][i] = (np.sqrt(G/l))/(2*np.pi)
    
        i = i + 1

    mp.plot(frequence[1],frequence[0])
    mp.plot(frequence[1],frequence[2],'r')
    mp.title('frequence en fonction de theta initial ')
    mp.xlabel('theta(0) ')
    mp.ylabel('frequence ')
    
    mp.savefig("frequences.png")                
    mp.show()





#freq(200)

 






#Modelisation du probleme du pendule a 2 maillons

def func2():

    def f(theta,t):

        return theta[2]
        
    def g(theta,t):

        return theta[3]

    def d(theta,t):

        return (-G*3*m*np.sin(theta[0])-m*G*np.sin(theta[0]-2*theta[1])-2*np.sin(theta[0]-theta[1])*m*(l*theta[3]*theta[3]+l*theta[2]*theta[2]*np.cos(theta[0]-theta[1]))) / (l*(3*m-m*np.cos(2*theta[0]-2*theta[1]))) 

    def j(theta,t):

        return (2*np.sin(theta[0]-theta[1])*(2*m*l*theta[2]*theta[2]+2*G*m*np.cos(theta[0])+theta[3]*theta[3]*l*m*np.cos(theta[0]-theta[1]))) / (l*(3*m-m*np.cos(2*theta[0]-2*theta[1])))


    return [f,g,d,j]


omega0 = func2()


#Resolution


#print "1 maillons:"
Y1 = pa.meth_step_n(np.array([(np.pi)/2,0.]),0.,50,step,theta0,pa.runge_kutta_step)[0]
#print Y1
#print "\n"


A = l*np.sin(Y1[:,0])
B = -l*np.cos(Y1[:,0])
             
#Trace du resultat


#mp.clf()


#mp.plot(A,B,'b')
#mp.plot(np.arange(0,21,1),Y1[:,3],'r')

#mp.show()
                         
print "2 maillons:"
Y2 = pa.meth_step_n(np.array([2.,-0.5,0.,0.]),0.,1000,step,omega0,pa.runge_kutta_step)[0]
print Y2
print "\n"


x1 = l*np.sin(Y2[:,0])
y1 = -l*np.cos(Y2[:,0])
x2 = l*np.sin(Y2[:,0]) + l*np.sin(Y2[:,1])
y2 = -l*np.cos(Y2[:,0]) + -l*np.cos(Y2[:,1])
             
#Trace du resultat

#mp.clf()
#mp.plot(x1,y1,'b')
#mp.plot(x2,y2,'r')
#mp.title('trajectoire du pendule a deux maillons')
#mp.xlabel('x(t) ')
#mp.ylabel('y(t) ')

#mp.savefig("deux-maillons2.png")                
#mp.show()


#premier retournement

def ret(Y,i):
    n = np.shape(Y)[0]
    k = 0
    while ((np.absolute(Y[:,i][k]<np.pi)) and (k<(n-2))):
        k = k+1
    return (k-1)*step

#print ret(Y2,1)

def carte():
    S = np.zeros([13,13])
    R = np.zeros([13,13])
    for k in np.arange(0,13,1):
        for j in np.arange(0,13,1):
            R[k][j] =ret( pa.meth_step_n(np.array([k*0.5-3,j*0.5-3,0.,0.]),0.,1000,step,omega0,pa.runge_kutta_step)[0],1)
            if (R[k][j]<10*np.sqrt(l/G)):
                S[k][j] = 1.0
            elif (R[k][j]<100*np.sqrt(l/G)):
                S[k][j] = 2.0
            elif (R[k][j]<1000*np.sqrt(l/G)):
                S[k][j] = 3.0
            else :
                S[k][j] = 5.0
    return S


#mp.clf()
#S = carte()
#mp.imshow(S,extent=(-3.,3.,-3.,3.))
#mp.title('premier retournement')
#mp.xlabel('theta1 initial ')
#mp.ylabel('theta2 initial ')

#mp.savefig("premier_retournement.png")
#mp.show()
